I tried to commit again, after fixing the problems found in the tests.

The function looks good overall, nice work! Although visualization functions can be hard to test, you should still add some basic tests into ./tests/utils/test_visualization.py. That file has some examples of how we test some of the other vis functions which you can follow. It's hard to test functionality with plotting functions, but just making sure the function runs with various input combinations should suffice.